{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Re-Identification using TAO ReIdentificationNet\n",
    "\n",
    "Transfer learning is the process of transferring learned features from one application to another. It is a commonly used training technique where you use a model trained on one task and re-train to use it on a different task. \n",
    "\n",
    "Train Adapt Optimize (TAO) Toolkit  is a simple and easy-to-use Python based AI toolkit for taking purpose-built AI models and customizing them with users' own data.\n",
    "\n",
    "<img align=\"center\" src=\"https://d29g4g2dyqv443.cloudfront.net/sites/default/files/akamai/TAO/tlt-tao-toolkit-bring-your-own-model-diagram.png\" width=\"1080\">\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Learning Objectives\n",
    "\n",
    "In this notebook, you will learn how to leverage the simplicity and convenience of TAO to:\n",
    "\n",
    "* Train a model for re-identification on the [Market-1501](https://zheng-lab.cecs.anu.edu.au/Project/project_reid.html) dataset.\n",
    "* Evaluate the trained model & visualize results.\n",
    "* Run Inference on the trained model.\n",
    "* Export the trained model to an ONNX file for deployment to DeepStream or TensorRT.\n",
    "\n",
    "At the end of this notebook, you will have generated a trained and optimized `reidentification` model, \n",
    "which you may deploy with this [end-to-end sample](https://github.com/NVIDIA-AI-IOT/tao-toolkit-triton-apps) with Triton.\n",
    "\n",
    "## Table of Contents\n",
    "\n",
    "This notebook shows an example usecase of ReIdentificationNet using Train Adapt Optimize (TAO) Toolkit.\n",
    "\n",
    "0. [Set up env variables and map drives](#head-0)\n",
    "1. [Installing the TAO launcher](#head-1)\n",
    "2. [Prepare dataset and pre-trained model](#head-2)\n",
    "3. [Provide training specification](#head-3)\n",
    "4. [Run TAO training](#head-4)\n",
    "5. [Evaluate trained models](#head-5)\n",
    "6. [Inferences](#head-6)\n",
    "7. [Deploy](#head-7)\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 0. Set up env variables and map drives <a class=\"anchor\" id=\"head-0\"></a>\n",
    "\n",
    "When using the purpose-built pretrained models from NGC, please make sure to set the `$KEY` environment variable to the key as mentioned in the model overview. Failing to do so, can lead to errors when trying to load them as pretrained models.\n",
    "\n",
    "The TAO launcher uses docker containers under the hood, and **for our data and results directory to be visible to the docker, they need to be mapped**. The launcher can be configured using the config file `~/.tao_mounts.json`. Apart from the mounts, you can also configure additional options like the Environment Variables and amount of Shared Memory available to the TAO launcher. <br>\n",
    "\n",
    "`IMPORTANT NOTE:` The code below creates a sample `~/.tao_mounts.json`  file. Here, we can map directories in which we save the data, specs, results and cache. You should configure it for your specific case so these directories are correctly visible to the docker container.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "# Please define this local project directory that needs to be mapped to the TAO docker session.\n",
    "%env LOCAL_PROJECT_DIR=/path/to/local/tao-experiments\n",
    "\n",
    "os.environ[\"HOST_DATA_DIR\"] = os.path.join(os.getenv(\"LOCAL_PROJECT_DIR\", os.getcwd()), \"data\", \"reidentificationnet\")\n",
    "os.environ[\"HOST_RESULTS_DIR\"] = os.path.join(os.getenv(\"LOCAL_PROJECT_DIR\", os.getcwd()), \"reidentificationnet\")\n",
    "os.environ[\"HOST_MODEL_DIR\"] = os.path.join(os.getenv(\"LOCAL_PROJECT_DIR\", os.getcwd()), \"data\", \"reidentificationnet\", \"model\")\n",
    "\n",
    "# Set this path if you don't run the notebook from the samples directory.\n",
    "# %env NOTEBOOK_ROOT=/path/to/local/tao-experiments/re_identification_net\n",
    "# The sample spec files are present in the same path as the downloaded samples.\n",
    "os.environ[\"HOST_SPECS_DIR\"] = os.path.join(\n",
    "    os.getenv(\"NOTEBOOK_ROOT\", os.getcwd()),\n",
    "    \"specs\"\n",
    ")\n",
    "\n",
    "# Set your encryption key, and use the same key for all commands\n",
    "%env KEY = nvidia_tao"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "! mkdir -p $HOST_DATA_DIR\n",
    "! mkdir -p $HOST_SPECS_DIR\n",
    "! mkdir -p $HOST_RESULTS_DIR\n",
    "! mkdir -p $HOST_MODEL_DIR"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Mapping up the local directories to the TAO docker.\n",
    "import json\n",
    "import os\n",
    "mounts_file = os.path.expanduser(\"~/.tao_mounts.json\")\n",
    "tlt_configs = {\n",
    "   \"Mounts\":[\n",
    "       # Mapping the data directory\n",
    "       {\n",
    "           \"source\": os.environ[\"LOCAL_PROJECT_DIR\"],\n",
    "           \"destination\": \"/workspace/tao-experiments\"\n",
    "       },\n",
    "       {\n",
    "           \"source\": os.environ[\"HOST_DATA_DIR\"],\n",
    "           \"destination\": \"/data\"\n",
    "       },\n",
    "       {\n",
    "           \"source\": os.environ[\"HOST_MODEL_DIR\"],\n",
    "           \"destination\": \"/model\"\n",
    "       },\n",
    "       {\n",
    "           \"source\": os.environ[\"HOST_SPECS_DIR\"],\n",
    "           \"destination\": \"/specs\"\n",
    "       },\n",
    "       {\n",
    "           \"source\": os.environ[\"HOST_RESULTS_DIR\"],\n",
    "           \"destination\": \"/results\"\n",
    "       }\n",
    "   ],\n",
    "   \"DockerOptions\": {\n",
    "        \"shm_size\": \"16G\",\n",
    "        \"ulimits\": {\n",
    "            \"memlock\": -1,\n",
    "            \"stack\": 67108864\n",
    "         }\n",
    "   }\n",
    "}\n",
    "# Writing the mounts file.\n",
    "with open(mounts_file, \"w\") as mfile:\n",
    "    json.dump(tlt_configs, mfile, indent=4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!cat ~/.tao_mounts.json"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1. Installing the TAO launcher <a class=\"anchor\" id=\"head-1\"></a>\n",
    "The TAO launcher is a python package distributed as a python wheel listed in PyPI. You may install the launcher by executing the following cell.\n",
    "\n",
    "Please note that TAO Toolkit recommends users to run the TAO launcher in a virtual env with python 3.6.9. You may follow the instruction in this [page](https://virtualenvwrapper.readthedocs.io/en/latest/install.html) to set up a python virtual env using the `virtualenv` and `virtualenvwrapper` packages. Once you have setup virtualenvwrapper, please set the version of python to be used in the virtual env by using the `VIRTUALENVWRAPPER_PYTHON` variable. You may do so by running\n",
    "\n",
    "```sh\n",
    "export VIRTUALENVWRAPPER_PYTHON=/path/to/bin/python3.x\n",
    "```\n",
    "where x >= 6 and <= 8\n",
    "\n",
    "We recommend performing this step first and then launching the notebook from the virtual environment. In addition to installing TAO python package, please make sure of the following software requirements:\n",
    "* python >=3.7, <=3.10.x\n",
    "* docker-ce > 19.03.5\n",
    "* docker-API 1.40\n",
    "* nvidia-container-toolkit > 1.3.0-1\n",
    "* nvidia-container-runtime > 3.4.0-1\n",
    "* nvidia-docker2 > 2.5.0-1\n",
    "* nvidia-driver > 455+\n",
    "\n",
    "Once you have installed the pre-requisites, please log in to the docker registry nvcr.io by following the command below\n",
    "\n",
    "```sh\n",
    "docker login nvcr.io\n",
    "```\n",
    "\n",
    "You will be triggered to enter a username and password. The username is `$oauthtoken` and the password is the API key generated from `ngc.nvidia.com`. Please follow the instructions in the [NGC setup guide](https://docs.nvidia.com/ngc/ngc-overview/index.html#generating-api-key) to generate your own API key.\n",
    "\n",
    "Please note that TAO Toolkit recommends users to run the TAO launcher in a virtual env with python >=3.6.9. You may follow the instruction in this [page](https://virtualenvwrapper.readthedocs.io/en/latest/install.html) to set up a python virtual env using the virtualenv and virtualenvwrapper packages."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# SKIP this step IF you have already installed the TAO launcher.\n",
    "!pip3 install nvidia-tao"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# View the versions of the TAO launcher\n",
    "!tao info"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2. Prepare dataset and pre-trained model <a class=\"anchor\" id=\"head-2\"></a>\n",
    " We will be using the [Market-1501](https://zheng-lab.cecs.anu.edu.au/Project/project_reid.html) dataset for the tutorial. Download the dataset [here](https://drive.google.com/file/d/1TwkgQcIa_EgRjVMPSbyEKtcfljqURrzi/view?usp=sharing) and extract it. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Download the dataset.\n",
    "!pip3 install -U gdown\n",
    "!gdown https://drive.google.com/uc?id=0B8-rUzbwVRk0c054eEozWG9COHM -O $HOST_DATA_DIR/market1501.zip"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Extract the files\n",
    "!unzip -u $HOST_DATA_DIR/market1501.zip -d $HOST_DATA_DIR\n",
    "!rm -rf $HOST_DATA_DIR/market1501\n",
    "!mv $HOST_DATA_DIR/Market-1501-v15.09.15 $HOST_DATA_DIR/market1501\n",
    "!rm $HOST_DATA_DIR/market1501.zip"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Verify\n",
    "!ls -l $HOST_DATA_DIR/market1501"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import glob\n",
    "import re\n",
    "import random\n",
    "\n",
    "def sample_dataset(input_dir, output_dir, n_samples, use_ids = None):\n",
    "    \"\"\"Select a subset of images fom input_dir and move them to output_dir.\n",
    "    \n",
    "    Args:\n",
    "        input_dir (str): Input Folder Path of the train images.\n",
    "        output_dir (str): Output Folder Path of the test images.\n",
    "        n_samples (int): Number of samples to use.\n",
    "        use_ids(list int): List of IDs to grab from test and query folder.\n",
    "        \n",
    "    Returns:\n",
    "        IDs used for sampling\n",
    "    \"\"\"\n",
    "    img_paths = glob.glob(os.path.join(input_dir, '*.jpg'))\n",
    "    pattern = re.compile(r'(\\d+)_c(\\d+)')\n",
    "    id_to_img = {}\n",
    "\n",
    "    # Grab images with matching ids\n",
    "    for img_path in img_paths:\n",
    "        pid, _ = map(int, pattern.search(img_path).groups())\n",
    "        if pid not in id_to_img:\n",
    "            id_to_img[pid] = []\n",
    "        id_to_img[pid].append(img_path)\n",
    "    \n",
    "    # Create directory\n",
    "    if not os.path.exists(output_dir):\n",
    "        os.makedirs(output_dir)\n",
    "    else:\n",
    "        command = \"rm -r \" + output_dir\n",
    "        os.system(command)\n",
    "        os.makedirs(output_dir)\n",
    "\n",
    "    assert id_to_img, \"Dataset size cannot be 0.\"\n",
    "\n",
    "    sampled_ids = random.sample(list(id_to_img.keys()), n_samples)\n",
    "\n",
    "    for sampled_id in sampled_ids:\n",
    "        for img_path in id_to_img[sampled_id]:\n",
    "            command = \"cp \" + img_path + \" \" + output_dir\n",
    "            os.system(command)\n",
    "\n",
    "    # Use same ids for test and query\n",
    "    if use_ids:    \n",
    "        \n",
    "        # Create query dir\n",
    "        if not os.path.exists(output_dir):\n",
    "            os.makedirs(output_dir)\n",
    "        else:\n",
    "            command = \"rm -r \" + output_dir\n",
    "            os.system(command)\n",
    "            os.makedirs(output_dir)\n",
    "\n",
    "        # Find images in test with same id\n",
    "        img_paths = glob.glob(os.path.join(input_dir, '*.jpg'))\n",
    "        for use_id in use_ids:\n",
    "            for img_path in img_paths:\n",
    "                person_id, _ = map(int, pattern.search(img_path).groups())\n",
    "                if use_id == person_id:\n",
    "                    command = \"cp \" + img_path + \" \" + output_dir\n",
    "                    os.system(command)\n",
    "\n",
    "    return sampled_ids\n",
    "\n",
    "# Number of samples\n",
    "n_samples = 100\n",
    "data_dir = os.path.join(os.environ[\"HOST_DATA_DIR\"], \"market1501\")\n",
    "\n",
    "# Create train dataset\n",
    "train_input_dir = os.path.join(data_dir, \"bounding_box_train\")\n",
    "train_output_dir = os.path.join(data_dir, \"sample_train\")\n",
    "sample_dataset(train_input_dir, train_output_dir, n_samples)\n",
    "\n",
    "# Create test dataset\n",
    "test_input_dir = os.path.join(data_dir, \"bounding_box_test\")\n",
    "test_output_dir = os.path.join(data_dir, \"sample_test\")\n",
    "ids = sample_dataset(test_input_dir, test_output_dir, n_samples)\n",
    "\n",
    "# Create query dataset\n",
    "query_input_dir = os.path.join(data_dir, \"query\")\n",
    "query_output_dir = os.path.join(data_dir, \"sample_query\")\n",
    "sample_dataset(query_input_dir, query_output_dir, n_samples, ids)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Verify the sample_* folder is not empty\n",
    "!ls -l $HOST_DATA_DIR/market1501"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# We provide pre-trained models for Swin-Tiny & Swin-Base variants.\n",
    "# Swin-Tiny offers faster training whereas Swin-Base offers higher accuracy.\n",
    "\n",
    "# For faster training, use the pre-trained Swin-Tiny model.\n",
    "# Download it using the following command:\n",
    "!mkdir -p $HOST_MODEL_DIR/market1501\n",
    "!gdown https://drive.google.com/uc?id=1sPLlpcqluyncqSnps88WfkfbshXyXB7F -O $HOST_MODEL_DIR/market1501/swintiny_pretrained.pth\n",
    "\n",
    "# For higher accuracy, use the pre-trained Swin-Base model.\n",
    "# Download it using the following command:\n",
    "# !gdown https://drive.google.com/uc?id=12R97_bSVLCKsPx1E8zPj0ltK6_MDYGhn -O $HOST_MODEL_DIR/market1501/swinbase_pretrained.pth"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3. Provide training specification <a class=\"anchor\" id=\"head-3\"></a>\n",
    "\n",
    "We provide specification files to configure the training parameters including:\n",
    "\n",
    "* model: configure the model setting\n",
    "  * backbone: type of backbone architecture, example: swin_tiny_patch4_window7_224, swin_base_patch4_window7_224\n",
    "  * last_stride: stride size\n",
    "  * pretrain_choice: type of pretrain choice, self\n",
    "  * pretrained_model_path: path for the input model\n",
    "  * input_channels: number of input channels\n",
    "  * input_width: width of an input image\n",
    "  * input_height: height of an input image\n",
    "  * neck: type of neck, bnneck\n",
    "  * stride_size: size of the stride\n",
    "  * reduce_feat_dim: reduce the dimension of the output embedding\n",
    "  * feat_dim: size of the output embedding\n",
    "  * no_margin: with soft triplet loss\n",
    "  * neck_feat: statergy to place neck, after/before\n",
    "  * metric_loss_type: loss type for model, triplet/center/triplet_center\n",
    "  * with_center_loss: enabling center loss feature, True/False\n",
    "  * with_flip_feature: enabling image flipping feature, True/False\n",
    "  * label_smooth: enabling label smoothing feature, True/False\n",
    "  * pretrain_hw_ratio: height-width ratio of pre-trained model\n",
    "* dataset: configure the dataset and augmentation methods\n",
    "  * train_dataset_dir: path for the train dataset directory\n",
    "  * test_dataset_dir: path for the test dataset directory\n",
    "  * query_dataset_dir: path for the query dataset directory\n",
    "  * num_classes: number of classes\n",
    "  * batch_size: number of images in 1 batch for training\n",
    "  * val_batch_size: number of images in 1 batch for validation\n",
    "  * num_workers: number of workers to do data loading\n",
    "  * pixel_mean: pixel mean in 3 channels for normalization\n",
    "  * pixel_std: pixel standard deviation in 3 channels for normalization\n",
    "  * padding: add padding to images\n",
    "  * prob: probability of randomly flipping images horizontally\n",
    "  * re_prob: constant for random erasing\n",
    "  * sampler: stratergy to load images, softmax_triplet\n",
    "  * num_instances: number of types 1 images is repeated in a batch\n",
    "* re_ranking: configure the re-ranking features\n",
    "  * re_ranking: enabling re-ranking feature, True/False\n",
    "  * k1: constant for re-reranking feature\n",
    "  * k2: constant for re-reranking feature\n",
    "  * lambda_value: constant for re-reranking feature\n",
    "* train: configure the training hyperparameters\n",
    "  * optim: configure optimizer\n",
    "  * num_epochs: number of epochs\n",
    "  * checkpoint_interval: enabling how often to store models\n",
    "\n",
    "Please refer to the TAO documentation about ReIdentificationNet to get all the parameters that are configurable."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "!cat $HOST_SPECS_DIR/experiment_market1501_swin.yaml"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 4. Run TAO training <a class=\"anchor\" id=\"head-4\"></a>\n",
    "* Provide the sample spec file and the output directory location for models.\n",
    "* WARNING: Training will take several hours or one day to complete."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# NOTE: The following paths are set from the perspective of the TAO Docker.\n",
    "\n",
    "# The data is saved here\n",
    "%env DATA_DIR = /data\n",
    "%env MODEL_DIR = /model\n",
    "%env SPECS_DIR = /specs\n",
    "%env RESULTS_DIR = /results"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 4.1 Train Market-1501 model"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We will train a Market-1501 model from scratch."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "print(\"Train model\")\n",
    "!tao model re_identification train \\\n",
    "                  -e $SPECS_DIR/experiment_market1501_swin.yaml \\\n",
    "                  results_dir=$RESULTS_DIR/market1501 \\\n",
    "                  encryption_key=$KEY"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print('Listing checkpoints:')\n",
    "print('---------------------')\n",
    "!ls -ltrh $HOST_RESULTS_DIR/market1501/train"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# You can set NUM_EPOCH to the epoch corresponding to any saved checkpoint\n",
    "# %env NUM_EPOCH=029\n",
    "\n",
    "# Get the name of the checkpoint corresponding to your set epoch\n",
    "# tmp=!ls $HOST_RESULTS_DIR/market1501/train/*.pth | grep epoch_$NUM_EPOCH\n",
    "# %env CHECKPOINT={tmp[0]}\n",
    "\n",
    "# Or get the latest checkpoint\n",
    "os.environ[\"CHECKPOINT\"] = os.path.join(os.getenv(\"HOST_RESULTS_DIR\"), \"market1501/train/reid_model_latest.pth\")\n",
    "\n",
    "print('Rename a trained model: ')\n",
    "print('---------------------')\n",
    "!cp $CHECKPOINT $HOST_RESULTS_DIR/market1501/train/swintiny_market1501_model.tlt\n",
    "!ls -ltrh $HOST_RESULTS_DIR/market1501/train/swintiny_market1501_model.tlt"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 5. Evaluate trained models <a class=\"anchor\" id=\"head-5\"></a>\n",
    "Evaluate trained model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!tao model re_identification evaluate \\\n",
    "                    -e $SPECS_DIR/experiment_market1501_swin.yaml \\\n",
    "                    results_dir=$RESULTS_DIR/market1501 \\\n",
    "                    encryption_key=$KEY \\\n",
    "                    evaluate.results_dir=$RESULTS_DIR/market1501/evaluate \\\n",
    "                    evaluate.checkpoint=$RESULTS_DIR/market1501/train/swintiny_market1501_model.tlt \\\n",
    "                    evaluate.output_sampled_matches_plot=$RESULTS_DIR/market1501/evaluate/sampled_matches.png \\\n",
    "                    evaluate.output_cmc_curve_plot=$RESULTS_DIR/market1501/evaluate/cmc_curve.png \\\n",
    "                    evaluate.test_dataset=$DATA_DIR/market1501/sample_test \\\n",
    "                    evaluate.query_dataset=$DATA_DIR/market1501/sample_query"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The sampled matches and Cumulative Matching Characteristics (CMC) curve output can be visualized using the following script."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import IPython\n",
    "from PIL import Image\n",
    "\n",
    "file_name = os.path.join(os.environ[\"HOST_RESULTS_DIR\"], \"market1501\", \"evaluate\",  \"sampled_matches.png\")\n",
    "IPython.display.display(Image.open(file_name))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "file_name = os.path.join(os.environ[\"HOST_RESULTS_DIR\"], \"market1501\", \"evaluate\",  \"cmc_curve.png\")\n",
    "IPython.display.display(Image.open(file_name))"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 6. Inferences <a class=\"anchor\" id=\"head-6\"></a>\n",
    "In this section, we run the re-identification inference tool to generate inferences with the trained models and save the results under `$RESULTS_DIR`. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!tao model re_identification inference \\\n",
    "                    -e $SPECS_DIR/experiment_market1501_swin.yaml \\\n",
    "                    results_dir=$RESULTS_DIR/market1501 \\\n",
    "                    encryption_key=$KEY \\\n",
    "                    inference.results_dir=$RESULTS_DIR/market1501/inference \\\n",
    "                    inference.checkpoint=$RESULTS_DIR/market1501/train/swintiny_market1501_model.tlt \\\n",
    "                    inference.output_file=$RESULTS_DIR/market1501/inference/inference.json \\\n",
    "                    inference.test_dataset=$DATA_DIR/market1501/sample_test \\\n",
    "                    inference.query_dataset=$DATA_DIR/market1501/sample_query"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 7. Deploy <a class=\"anchor\" id=\"head-7\"></a>\n",
    "Export the model to an ONNX model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "!tao model re_identification export \\\n",
    "                   -e $SPECS_DIR/experiment_market1501_swin.yaml \\\n",
    "                   results_dir=$RESULTS_DIR/market1501 \\\n",
    "                   encryption_key=$KEY \\\n",
    "                   export.checkpoint=$RESULTS_DIR/market1501/train/swintiny_market1501_model.tlt \\\n",
    "                   export.onnx_file=$RESULTS_DIR/market1501/export/swintiny_market1501_model.onnx"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print('Exported model:')\n",
    "print('------------')\n",
    "!ls -lth $HOST_RESULTS_DIR/market1501/export"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You may continue by deploying the exported model to [Triton Inference Server](https://developer.nvidia.com/nvidia-triton-inference-server). Please refer to the [TAO Toolkit Triton Apps](https://github.com/NVIDIA-AI-IOT/tao-toolkit-triton-apps), where a sample for end-to-end inference is provided. "
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.8.10 64-bit",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.10"
  },
  "vscode": {
   "interpreter": {
    "hash": "767d51c1340bd893661ea55ea3124f6de3c7a262a8b4abca0554b478b1e2ff90"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
